"""
这个正逆解是针对于垃圾工装弄的，它的x轴向大概为21.3mm，z轴向大概为203.85mm
正解不必赘述，顺着建立好的坐标系求就完事了
逆解大概的思路是需要A6轴的配合，A6轴设为0我暂时想不到怎么求解，大概如果用A6轴得他妈的左面五个矩阵乘法，右面一个矩阵乘法，还都带未知数，也许会更难，总之哥们放弃了
大概是已知了工具坐标系P'点沿着法向量-F向上走了203.85mm，求得P点的坐标，法向量沿用原来的，只不过全部取反变为F
此时因为我的探头朝下，所以它的法向量F的z必然为负，此时令y=0，则根据x^2+z^2=r^2，且x*Fx+z*Fz=0，解得A4,A5,A6轴的交汇点
交汇点Q就是P点加上上面求得x和z，此时利用之前的逆解求出A1,A2,A3的角度，此时再令A6=0，根据法向量F求得A4,A5的角度
关键点来了，我要在当前坐标系下，求得6号坐标系的X轴位置，然后计算它和向量QP，依此来确定A6轴的角度
理论成立，开始代码！！！
"""

import numpy as np

def dh_transform_modified(theta, d, a, alpha):
    """
    计算旋转矩阵
    """
    theta = np.radians(theta)
    alpha = np.radians(alpha)
    return np.array([
        [np.cos(theta), -np.sin(theta), 0, a],
        [np.sin(theta) * np.cos(alpha), np.cos(theta) * np.cos(alpha), -np.sin(alpha), -d * np.sin(alpha)],
        [np.sin(theta) * np.sin(alpha), np.cos(theta) * np.sin(alpha), np.cos(alpha), d * np.cos(alpha)],
        [0, 0, 0, 1]
    ])

def forward_kinematics(A1, A2, A3, A4, A5, A6,XTool=21.3,YTool=0,ZTool=203.85):
    # 定义DH参数
    dh_params = [
        [-A1, 400, 0, 0],
        [A2, 0, 25, -90],
        [A3 - 90, 0, 455, 0],
        [-A4, 420, 25, -90],
        [A5, 0, 0, 90],
        [- A6, 0, 0, -90]
    ]

    # 计算6号坐标系相对于基座标的变换矩阵
    T = np.eye(4)  # 做个单位矩阵
    for theta, d, a, alpha in dh_params:
        T = T @ dh_transform_modified(theta, d, a, alpha)

    # 计算探头在基坐标系的位置
    tool_coords = np.array([XTool, YTool, ZTool, 1])
    base_coords = T @ tool_coords

    # Extract x, y, z from the resulting matrix
    x, y, z = base_coords[0], base_coords[1], base_coords[2]
    # 打印结果
    # print(f"x: {x}, y: {y}, z: {z}")
    coordinationList=[x,y,z]
    return coordinationList


def inverse_kinematics(x,y,z,px,py,pz,toolSizeX=21.3,toolSizeZ=203.85):
    """
    求解A1, A2, A3的角度。
    :param P3: 在n=3坐标系中的点 (x3, y3, z3)
    :param d1: d1参数 (单位: mm)
    :param a2: a2参数 (单位: mm)
    :param a3: a3参数 (单位: mm)
    :return: A1, A2, A3 (单位: 度)
    """
    # 计算点P
    P1 = np.array([-px, -py, -pz])
    P1 = P1 / np.linalg.norm(P1)
    px, py, pz = -P1[0], -P1[1], -P1[2]
    x=x+px*toolSizeZ
    y=y+py*toolSizeZ
    z=z+pz*toolSizeZ
    # 计算点Q
    if px==0 and py==0:
        zIncrease=0
        yIncrease=0
        xIncrease=toolSizeX
    elif py==0:
        yIncrease = 0
        zIncrease=toolSizeX*px*np.sqrt(1/(pz**2+px**2))
        xIncrease=-pz*zIncrease/px
    elif px==0:
        xIncrease = 0
        zIncrease=toolSizeX*py*np.sqrt(1/(pz**2+py**2))
        yIncrease=-pz*zIncrease/py
    else:
        yIncrease = 0
        zIncrease = toolSizeX * px * np.sqrt(1 / (pz ** 2 + px ** 2))
        xIncrease = -pz * zIncrease / px
    x0=x+xIncrease
    y0=y+yIncrease
    z0=z+zIncrease
    # 打印结果
    # print(f"x0: {x0}, y0: {y0}, z0: {z0}")


    A1=np.arctan2(y0, x0)
    a=np.sqrt(x0**2+y0**2)
    b=z0-400
    d=np.sqrt((a-25)**2+b**2)
    a3_2=np.sqrt(420**2+25**2)
    if d>a3_2+455 or d<455-a3_2:
        raise ValueError("点超出了机械臂的工作空间！")
    theta1=np.arctan2(b,a-25)
    theta2=np.arccos((455**2+((a-25)**2+b**2)-(420**2+25**2))/(2*455*d))
    A2=theta1+theta2 #一会要注意变成负数
    A3=np.pi-np.arccos((-((a-25)**2+b**2)+(420**2+25**2)+455**2)/(2*a3_2*455))+np.arctan2(25,420)

    # 此时已求得A1,A2,A3的角度，接下来从旋转矩阵的坐标系角度开始求解A4,A5,A6的角度

    # x_rela是关节轴A3的位置
    x_rela=25+455*np.cos(A2)
    y_rela=400+455*np.sin(A2)

    vec1_x=a-x_rela
    vec1_y=z0-y_rela
    theta3=np.arctan2(vec1_y,vec1_x)-np.arctan2(25,420)
    # theta3是指在抛出A1影响下的平面上，此时Z4轴的角度
    vec_z=np.array([np.cos(theta3)*np.cos(A1),np.cos(theta3)*np.sin(A1),np.sin(theta3)])
    normalized_vec_z = vec_z / np.linalg.norm(vec_z)

    theta4=theta3+np.pi/2 #theta4是指的4号坐标系在A4=0时的X轴的位置
    vec_x=np.array([np.cos(theta4)*np.cos(A1),np.cos(theta4)*np.sin(A1),np.sin(theta4)])
    normalized_vec_x = vec_x / np.linalg.norm(vec_x)
    normalized_vec_y=np.array([np.sin(A1),-np.cos(A1),0])
    rotation_matrix=np.array([normalized_vec_x,normalized_vec_y,normalized_vec_z])
    # rotation_matrix代表的是旋转坐标系，它的转置是整个4号坐标系在A4=0时的相对于基坐标系的旋转矩阵，也可以代表4号坐标系各个轴的指向，两者等价
    # 因为旋转矩阵的特性，我要求P1在4号坐标系下的坐标需要求得0号坐标系相对于四号坐标系的旋转矩阵，而这恰好是rotaion_matrix的转置，所以并不需要求转置
    P1_rotated=np.dot(rotation_matrix,P1)# 代表了4号坐标系下的探头法向量指向的各个轴分量
    A4=np.pi-np.arctan2(P1_rotated[1],P1_rotated[0])
    # A4=np.arctan2(P1_rotated[0],P1_rotated[1])
    A5=np.arccos(P1_rotated[2])
    # 后面就自动默认A6=0了
    # 计算6号坐标系的X轴位置，需要计算旋转矩阵R4'4，R54，R65而R04是已知的，即rotation_matrix
    R4_Four=np.array([np.array([np.cos(A4),-np.sin(A4),0]),np.array([np.sin(A4),np.cos(A4),0]),np.array([0,0,1])]).T
    R5_4=np.array([np.array([np.cos(A5),0,np.sin(A5)]),np.array([np.sin(-A5),0,np.cos(A5)]),np.array([0,-1,0])]).T
    R6_5=np.array([np.array([1,0,0]),np.array([0,0,-1]),np.array([0,1,0])]).T
    # print(f"R4_Four: {R4_Four}")
    # print(f"R6_5: {R6_5}")
    # print(f"测试矩阵乘法：{np.dot(R6_5,np.array([3,2,1]))}")
    R6_0=np.dot(np.dot(np.dot(rotation_matrix.T,R4_Four),R5_4),R6_5)

    # R6x代表了6号坐标系的X轴指向的各个轴分量，即为我此时如果A6轴为0时的腕关节对探头反向延长线交汇点的向量在基座标的方向
    # Probe_Real代表了基座标下腕关节与探头两点的向量,这三个分量现在都是单位向量，可以直接用numpy里的反三角函数
    R6x = np.dot(R6_0, np.array([1, 0, 0]))
    R6z = np.dot(R6_0, np.array([0, 0, 1]))
    Probe_Real=np.array([-xIncrease,-yIncrease,-zIncrease])/np.linalg.norm([-xIncrease,-yIncrease,-zIncrease])

    # 在以R6z为转轴的情况下，A6轴旋转使得R6x顺时针旋转
    # 下面通过向量叉乘确定A6轴的角度
    CrossVirtualToReal=np.cross(R6x,Probe_Real)
    dotVirtualToReal=np.dot(R6x,Probe_Real)
    A6temp=np.arctan2(np.linalg.norm(CrossVirtualToReal),dotVirtualToReal)
    if np.dot(R6z,CrossVirtualToReal)<0:
        A6=A6temp
    else:
        A6=-A6temp


    # 转换为角度
    A1=-A1
    A1_deg = np.degrees(A1).astype(np.float32)
    A2 = -A2
    A2_deg = np.degrees(A2).astype(np.float32)
    A3_deg = np.degrees(A3).astype(np.float32)
    A4_deg = np.degrees(A4).astype(np.float32)
    if A4_deg>181:
        A4_deg=A4_deg-360
    elif A4_deg<-181:
        A4_deg = 360 + A4_deg
    A5_deg = np.degrees(A5).astype(np.float32)
    A6_deg = np.degrees(A6).astype(np.float32)

    if A1_deg>170 or A1_deg<-170:
        raise ValueError("A1角度超出了机器限位角度！")
    if A2_deg>45 or A2_deg<-190:
        raise ValueError("A2角度超出了机器限位角度！")
    if A3_deg>156 or A3_deg<-120:
        raise ValueError("A3角度超出了机器限位角度！")
    if A4_deg>185 or A4_deg<-185:
        raise ValueError("A4角度超出了机器限位角度！")
    if A5_deg>120 or A5_deg<-120:
        raise ValueError("A5角度超出了机器限位角度！")
    # # 打印结果
    # print(f"A1: {A1_deg}°")
    # print(f"A2: {A2_deg}°")
    # print(f"A3: {A3_deg}°")
    # print(f"A4: {A4_deg}°")
    # print(f"A5: {A5_deg}°")
    # print(f"A6: {A6_deg}°")
    # print(f"R6x: {R6x}")
    # print(f"R6z: {R6z}")
    return A1_deg, A2_deg, A3_deg, A4_deg, A5_deg, A6_deg

# # 测试代码
# P3 = (75, 203, 756)
# P4 = (7, 9, 11)
# A1, A2, A3 =P3[0],P3[1],P3[2]
# A4, A5, A6 =P4[0],P4[1],P4[2]
# h,i,j,k,l,m=inverse_kinematics(A1, A2, A3, A4, A5, A6)
# forward_kinematics(h,i,j,k,l,m)

